import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from tkinter import *
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib

matplotlib.use("TkAgg")
import torch
from segment_anything import sam_model_registry, SamPredictor
from tqdm import tqdm


class ClothingColorExtractor:
    def __init__(self, root):
        self.root = root
        self.root.title("服装主色提取系统")
        self.root.geometry("1000x700")

        # 确保中文显示正常
        plt.rcParams["font.family"] = ["SimHei"]

        # 创建界面
        self.create_widgets()

        # 数据存储
        self.selected_image_path = None
        self.processed_images = {}
        self.test_report = []

        # 初始化SAM模型
        sam_checkpoint = "sam_vit_h_4b8939.pth"
        device = "cuda" if torch.cuda.is_available() else "cpu"
        sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
        sam.to(device=device)
        self.predictor = SamPredictor(sam)

    def create_widgets(self):
        main_frame = Frame(self.root)
        main_frame.pack(fill=BOTH, expand=True, padx=10, pady=10)

        # 左侧控制面板
        control_frame = Frame(main_frame, width=250)
        control_frame.pack(side=LEFT, fill=Y, padx=5, pady=5)

        path_frame = LabelFrame(control_frame, text="路径设置")
        path_frame.pack(fill=X, padx=5, pady=5)

        Label(path_frame, text="数据集路径:").grid(row=0, column=0, sticky=W, padx=5, pady=5)
        self.dataset_path_var = StringVar(value=r"img/images")
        self.dataset_path_entry = Entry(path_frame, textvariable=self.dataset_path_var, width=20)
        self.dataset_path_entry.grid(row=0, column=1, padx=5, pady=5)
        Button(path_frame, text="浏览", command=self.browse_dataset_path).grid(row=0, column=2, padx=5, pady=5)

        Button(control_frame, text="处理数据集", command=self.process_dataset, height=2).pack(fill=X, padx=5, pady=10)
        Button(control_frame, text="选择图片", command=self.select_image, height=2).pack(fill=X, padx=5, pady=5)
        Button(control_frame, text="提取主色", command=self.extract_color, height=2).pack(fill=X, padx=5, pady=5)

        Button(control_frame, text="显示颜色标记", command=self.display_color_masks, height=2).pack(fill=X, padx=5, pady=5)

        cluster_frame = LabelFrame(control_frame, text="聚类参数")
        cluster_frame.pack(fill=X, padx=5, pady=5)

        Label(cluster_frame, text="主色数量:").grid(row=0, column=0, sticky=W, padx=5, pady=5)
        self.color_num_var = IntVar(value=3)
        Spinbox(cluster_frame, from_=1, to=10, textvariable=self.color_num_var, width=5).grid(row=0, column=1, padx=5, pady=5)

        # 右侧显示区域
        display_frame = Frame(main_frame)
        display_frame.pack(side=RIGHT, fill=BOTH, expand=True, padx=5, pady=5)

        self.image_frame = LabelFrame(display_frame, text="图片")
        self.image_frame.pack(fill=BOTH, expand=True, padx=5, pady=5)

        self.image_label = Label(self.image_frame)
        self.image_label.pack(fill=BOTH, expand=True, padx=5, pady=5)

        self.result_frame = LabelFrame(display_frame, text="主色分析结果")
        self.result_frame.pack(fill=BOTH, expand=True, padx=5, pady=5)

        self.fig, (self.ax1, self.ax2, self.ax3) = plt.subplots(1, 3, figsize=(15, 5))
        self.fig.subplots_adjust(left=0.05, right=0.95, top=0.9, bottom=0.1, wspace=0.3)
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.result_frame)
        self.canvas.get_tk_widget().pack(fill=BOTH, expand=True)

    def browse_dataset_path(self):
        path = filedialog.askdirectory(title="选择数据集文件夹")
        if path:
            self.dataset_path_var.set(path)

    def select_image(self):
        path = filedialog.askopenfilename(title="选择图片", filetypes=[("图片文件", "*.jpg;*.jpeg;*.png;*.bmp")])
        if path:
            self.selected_image_path = path
            self.display_image(path)

    def display_image(self, path):
        img = Image.open(path)
        img.thumbnail((500, 500))
        photo = ImageTk.PhotoImage(img)
        self.image_label.config(image=photo)
        self.image_label.image = photo

    def process_dataset(self):
        dataset_path = self.dataset_path_var.get()
        if not os.path.exists(dataset_path):
            messagebox.showerror("错误", f"路径不存在: {dataset_path}")
            return

        image_files = []
        for file in os.listdir(dataset_path):
            if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                image_files.append(os.path.join(dataset_path, file))

        if not image_files:
            messagebox.showinfo("提示", f"路径下没有找到图片文件: {dataset_path}")
            return

        total = len(image_files)
        self.test_report = []
        for i, img_path in tqdm(enumerate(image_files), total=total, desc="处理图片"):
            try:
                result = self.analyze_image(img_path)
                self.processed_images[img_path] = result
                self.test_report.append({
                    "image_path": img_path,
                    "dominant_colors": result["colors"],
                    "percentages": result["percentages"]
                })
                print(f"已处理 {i + 1}/{total}: {img_path}")
            except Exception as e:
                print(f"处理失败 {img_path}: {str(e)}")

        self.generate_test_report()
        messagebox.showinfo("完成", f"已完成处理 {len(self.processed_images)} 张图片，测试报告已生成")

    def generate_test_report(self):
        report_path = "test_report.txt"
        with open(report_path, "w", encoding="utf-8") as f:
            f.write("数据集批量测试报告\n")
            f.write("=" * 50 + "\n")
            for item in self.test_report:
                f.write(f"图片路径: {item['image_path']}\n")
                f.write("主色及百分比:\n")
                for color, percentage in zip(item["dominant_colors"], item["percentages"]):
                    f.write(f"  颜色: {color}, 百分比: {percentage * 100:.2f}%\n")
                f.write("-" * 50 + "\n")

    def extract_color(self):
        if not self.selected_image_path:
            messagebox.showinfo("提示", "请先选择一张图片")
            return

        result = self.analyze_image(self.selected_image_path)
        self.processed_images[self.selected_image_path] = result
        self.display_result(result)

    def analyze_image(self, img_path):
        img = cv2.imread(img_path)
        if img is None:
            raise Exception(f"无法读取图片: {img_path}")

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        masked_img = self.clothing_detection_sam(img)
        colors, percentages = self.extract_dominant_colors(masked_img, self.color_num_var.get())

        return {
            'original_img': img,
            'masked_img': masked_img,
            'colors': colors,
            'percentages': percentages
        }

    def clothing_detection_sam(self, img):
        self.predictor.set_image(img)
        h, w = img.shape[:2]
        input_point = np.array([[w // 2, h // 2]])
        input_label = np.array([1])
        masks, scores, logits = self.predictor.predict(
            point_coords=input_point,
            point_labels=input_label,
            multimask_output=True,
        )
        mask = masks[np.argmax(scores)]
        masked_img = img * mask[:, :, np.newaxis]
        return masked_img

    def extract_dominant_colors(self, img, num_colors=3):
        pixels = img.reshape(-1, 3)
        pixels = pixels[~np.all(pixels == 0, axis=1)]

        if len(pixels) == 0:
            return [(0, 0, 0)], [1.0]

        kmeans = KMeans(n_clusters=num_colors, random_state=42)
        kmeans.fit(pixels)

        colors = kmeans.cluster_centers_.astype(int)
        labels = kmeans.labels_
        percentages = np.bincount(labels) / len(labels)

        sorted_indices = np.argsort(percentages)[::-1]
        colors = colors[sorted_indices]
        percentages = percentages[sorted_indices]

        return colors, percentages

    def display_result(self, result):
        self.ax1.clear()
        self.ax2.clear()
        self.ax3.clear()

        self.ax1.imshow(result['original_img'])
        self.ax1.set_title("原始图片")
        self.ax1.axis('off')

        colors = [tuple(c / 255) for c in result['colors']]
        labels = [f"{c}\n{int(p * 100)}%" for c, p in zip(result['colors'], result['percentages'])]

        wedge_props = {'edgecolor': 'black', 'linewidth': 1}
        self.ax2.pie(result['percentages'], labels=None, colors=colors, wedgeprops=wedge_props)
        self.ax2.set_title("服装主色分析")

        self.ax3.imshow(result['original_img'])
        self.ax3.set_title("颜色区域标记")
        self.ax3.axis('off')

        for i, (color, percentage) in enumerate(zip(result['colors'], result['percentages'])):
            color_mask = self.create_color_mask(result['masked_img'], color)
            rgba = np.zeros((color_mask.shape[0], color_mask.shape[1], 4))
            rgba[color_mask] = np.append(color / 255, 0.5)
            self.ax3.imshow(rgba, alpha=0.5)

        self.canvas.draw()

    def create_color_mask(self, img, target_color, threshold=50):
        img = img.astype(int)
        target_color = np.array(target_color, dtype=int)
        diff = np.abs(img - target_color)
        mask = np.all(diff <= threshold, axis=2)
        return mask

    def display_color_masks(self):
        if not self.selected_image_path:
            messagebox.showinfo("提示", "请先选择一张图片")
            return

        if self.selected_image_path in self.processed_images:
            result = self.processed_images[self.selected_image_path]
            self.display_result(result)
        else:
            result = self.analyze_image(self.selected_image_path)
            self.processed_images[self.selected_image_path] = result
            self.display_result(result)

if __name__ == "__main__":
    try:
        root = Tk()
        app = ClothingColorExtractor(root)
        root.mainloop()
    except KeyboardInterrupt:
        print("程序被手动中断，退出程序。")